package com.bsj.media.media;

import lombok.extern.slf4j.Slf4j;

import java.util.Arrays;

/**
 * <p>
 * Title: 解析H264数据
 * </p>
 * <p>
 * Description:
 * </p>
 * <p>
 * Company: BAJ
 * </p>
 *
 * @author JIANG
 * @date 2017-12-18
 */
@Slf4j
public class ParserH264 {


    static final int RET_OK = 0;
    static final int RET_ERR_FILE_NOT_EXIST = -1;
    static final int RET_ERR_UNKNOWN = -2;
    static final int RET_ERR_FILE_END = -3;
    static final int RET_ERR = -4;
    static final int RET_ERR_END_NAL_NOT_FOUND = -5;
    static final int RET_ERR_START_CODE_NOT_FOUND = -6;

    public static int ReadOneNaluFromBuf(byte[] buffer, int nBufferSize, int offSet, NalInfo pNalu) {
        int startPos = 0, endPos = 0;
        int remain = 0;
        byte nalFirstByte;
        int start0_code_start_size = 0;

        if (offSet + 3 >= nBufferSize) {
            return RET_ERR_START_CODE_NOT_FOUND;
        }

        startPos = FindStartCodeEx(buffer, offSet, nBufferSize - offSet);
        if (startPos < 0) {
            return RET_ERR_START_CODE_NOT_FOUND;
        }

        start0_code_start_size = (buffer[startPos + 2] == 0x1) ? 3 : 4;
        remain = nBufferSize - (startPos + start0_code_start_size);
        if (remain < 3) {
            return RET_ERR_END_NAL_NOT_FOUND;
        }

        endPos = FindStartCodeEx(buffer, startPos + start0_code_start_size, remain);
        if (endPos < 0) {
            return RET_ERR_END_NAL_NOT_FOUND;
        }

        pNalu.len = endPos - startPos; // 包括00 00 00 01
        pNalu.buf = Arrays.copyOfRange(buffer, startPos, endPos);
        nalFirstByte = (pNalu.buf[2] == 0x1) ? pNalu.buf[3] : pNalu.buf[4];
        pNalu.forbidden_bit = (nalFirstByte & 0x80) > 0 ? true : false;
        pNalu.nal_reference_idc = nalFirstByte & 0x60; // 2 bit
        pNalu.nal_unit_type = nalFirstByte & 0x1f; // 5 bit
        return endPos;
    }

    public static int FindStartCodeEx(byte[] buf, int offset, int size) {
        boolean bFound = false;
        int pos = offset;
        while (pos + 4 < buf.length) {
            if (buf[pos] != 0x0 || buf[pos + 1] != 0x0) {
                pos++;
                continue;
            }

            bFound = (buf[pos + 2] == 0x1) || (buf[pos + 2] == 0x0 && buf[pos + 3] == 0x1);
            if (bFound) {
                return pos;
            }
            pos += 1;
        }

        return -1;
    }

    /**
     * 读取BP帧的开始位置
     *
     * @param buf：外部数据
     * @param buf_size：外部数据大小
     * @return 成功则返回当前位置 , 失败则返回起始位置
     */
    public static int ReadBPFrameFromBuf(byte[] buf, int start_size, int buf_size) {
        if (buf == null) {
            return start_size;
        }

        for (int i = start_size; i < buf_size - 4; i++) {
            if ((buf[i] != 0x00) || (buf[i + 1] != 0x00)) {
                continue;
            }

            if (buf[i + 2] == 0x01) { // 000001
                return i;
            } else if ((buf[i + 2] == 0x00) && (buf[i + 3] == 0x01)) { // 00000001
                return i;
            } else {
                continue;
            }
        }

        return start_size;
    }

    /**
     * 获取I帧中的sps以及pps数据
     *
     * @param srcData   h264数据
     * @param srcOffset 偏移位
     * @param sps       sps系列参数集数据帧 sequence parameter set
     * @param pps       图像参数集合 picture parameter set
     * @return 返回I帧0x65起始位置, 去掉了h264帧头
     */
    public static int ReadIFrameSPSPPSFromBuf(byte[] srcData, int srcOffset, NalInfo sps, NalInfo pps) {

        int mediaLen = 0;
        // int preTagSize = 0;
        // h264 nalu帧
        NalInfo sei = new NalInfo();
        NalInfo nalu = new NalInfo();

        // 读取一帧数据 sps pps
        int ret = ParserH264.ReadOneNaluFromBuf(srcData, srcData.length, srcOffset, sps);
        if (ret == -6) {
            return -9;
        } else if (sps.getNal_unit_type() != 0x7) {
            return -7;
        }
        srcOffset = ret;
        ret = ParserH264.ReadOneNaluFromBuf(srcData, srcData.length, srcOffset, pps);
        if (ret < 0 || pps.getNal_unit_type() != 0x8) {
            // 走到这概率很小
            return -8;
        }
        srcOffset = ret;

        // 读取I帧数据，如果是SEI则剔除
        ret = ParserH264.ReadOneNaluFromBuf(srcData, srcData.length, srcOffset, sei);

        int naluLen = 0; // nalu 长度
        if (ret != -6) {
            if (ret > 0 && sei.getNal_unit_type() == 6) {
                // 这是SEI帧数据，不用，取后面的65
                srcOffset = ret;
                ret = ParserH264.ReadOneNaluFromBuf(srcData, srcData.length, srcOffset, nalu);
                if (ret == -5) {
                    // 表示只有h264 帧头，00 00 00 01，这是正常数据
                    ret = srcOffset;
                    // 4 为start code长度
                    naluLen = srcData.length - srcOffset - 4;
                } else if (ret == -6) {
                    // 表示没有h264帧头，异常数据
                    log.warn("关键帧中SEI帧后面没有0x65数据 异常：" + srcData.length);
                    mediaLen = -5;
                } else if (ret > 0 && nalu.getNal_unit_type() == 5) {
                    // 也是异常数据，说明 I 帧后面还有其他数据
                    log.warn("关键帧数据后面还有其他帧 异常：" + srcData.length);
                    naluLen = ret - srcOffset - 4;
                }
            } else if (ret == -5) {
                // 没有SEI帧
                naluLen = srcData.length - srcOffset - 4;
            }

            if (naluLen > 0) {
                // I 帧  srcOffset 为00 00 00 01 或 00 00 01起点

                // 实际需要发送数据长度
                mediaLen = srcOffset + 4;
            }
        } else {
            mediaLen = -5;
        }

        return mediaLen;
    }
}
